package scau.mis4.chanyulin;

import scau.mis4.chanyulin.base.Bits;
import scau.mis4.chanyulin.base.GF28;
import scau.mis4.chanyulin.utils.BitsUtils;
import scau.mis4.chanyulin.utils.ByteUtils;
import scau.mis4.chanyulin.utils.GF28Utils;
import scau.mis4.chanyulin.utils.HexUtils;
import java.util.*;

/**
 * Created by chanyulin on 2017/12/7.
 * 128位的AES加密实现（书本P114例子）
 */
public class AES {

    private static final String P_STR = "0123456789abcdeffedcba9876543210";    // 明文
    private static final String K_STR = "0f1571c947d9e8590cb7add6af7f6798";    // 密钥

    public static void main(String[] args) {

        // 将明文和密钥转换为字节
        byte[] pBytes = hexStr2AESBytes(P_STR);
        byte[] kBytes = hexStr2AESBytes(K_STR);
        // 初始状态
        printHexStrResult("【初始化】明文：", pBytes);
        printHexStrResult("【初始化】密钥：", kBytes);
        // 密钥扩展
        byte[] keyExpandRes = keyExpand(kBytes);
        // 第0轮的轮密钥加
        byte[] firstPlusRes = pkPlus(pBytes, kBytes);
        printHexStrResult("【初始化】轮密钥加的结果为：", firstPlusRes);
        // 10轮加密
        for (int i = 1; i <= 10; i++) {
            // 字节代替（S盒）
            byte[] sBoxRes = sBox(firstPlusRes);
            printHexStrResult("【第" + i + "轮】字节代替的结果为：", sBoxRes);
            // 行移位变换
            byte[] movRowsRes = movRows(sBoxRes);
            printHexStrResult("【第" + i + "轮】行移位变换的结果为：", movRowsRes);
            // 列混淆
            byte[] rowsConfuseRes;
            if (i != 10) {
                rowsConfuseRes = rowsConfuse(movRowsRes);
                printHexStrResult("【第" + i + "轮】列混淆的结果为：", rowsConfuseRes);
            } else {
                rowsConfuseRes = movRowsRes;
            }
            // 生成轮密钥
            byte[] roundKeyRes = getExpandedKey(keyExpandRes, i);
            printHexStrResult("【第" + i + "轮】轮密钥为：", roundKeyRes);
            // 轮密钥加
            byte[] plusRes = pkPlus(rowsConfuseRes, roundKeyRes);
            printHexStrResult("【第" + i + "轮】轮密钥加的结果为：", plusRes);
            firstPlusRes = plusRes;
        }
        printHexStrResult("【结果】密文：", firstPlusRes);
    }

    /**
     * 列混淆
     * @param input
     * @return
     */
    public static byte[] rowsConfuse(byte[] input) {
        GF28[] result = new GF28[input.length];
        GF28[] nums = new GF28[4];
        for (int i = 0; i < 4; i++) {
            nums[0] = GF28Utils.multi(new GF28(input[i]), new GF28(new byte[] {0, 0, 0, 0, 0, 0, 0, 1, 0}));
            nums[1] = GF28Utils.multi(new GF28(input[4 + i]), new GF28(new byte[] {0, 0, 0, 0, 0, 0, 0, 1, 1}));
            nums[2] = new GF28(input[8 + i]);
            nums[3] = new GF28(input[12 + i]);
            result[i] = GF28Utils.xor(nums[3], GF28Utils.xor(nums[2], GF28Utils.xor(nums[1], nums[0])));
        }
        for (int i = 0; i < 4; i++) {
            nums[0] = new GF28(input[i]);
            nums[1] = GF28Utils.multi(new GF28(input[4 + i]), new GF28(new byte[] {0, 0, 0, 0, 0, 0, 0, 1, 0}));
            nums[2] = GF28Utils.multi(new GF28(input[8 + i]), new GF28(new byte[] {0, 0, 0, 0, 0, 0, 0, 1, 1}));
            nums[3] = new GF28(input[12 + i]);
            result[4 + i] = GF28Utils.xor(nums[3], GF28Utils.xor(nums[2], GF28Utils.xor(nums[1], nums[0])));
        }
        for (int i = 0; i < 4; i++) {
            nums[0] = new GF28(input[i]);
            nums[1] = new GF28(input[4 + i]);
            nums[2] = GF28Utils.multi(new GF28(input[8 + i]), new GF28(new byte[] {0, 0, 0, 0, 0, 0, 0, 1, 0}));
            nums[3] = GF28Utils.multi(new GF28(input[12 + i]), new GF28(new byte[] {0, 0, 0, 0, 0, 0, 0, 1, 1}));
            result[8 + i] = GF28Utils.xor(nums[3], GF28Utils.xor(nums[2], GF28Utils.xor(nums[1], nums[0])));
        }
        for (int i = 0; i < 4; i++) {
            nums[0] = GF28Utils.multi(new GF28(input[i]), new GF28(new byte[] {0, 0, 0, 0, 0, 0, 0, 1, 1}));
            nums[1] = new GF28(input[4 + i]);
            nums[2] = new GF28(input[8 + i]);
            nums[3] = GF28Utils.multi(new GF28(input[12 + i]), new GF28(new byte[] {0, 0, 0, 0, 0, 0, 0, 1, 0}));
            result[12 + i] = GF28Utils.xor(nums[3], GF28Utils.xor(nums[2], GF28Utils.xor(nums[1], nums[0])));
        }
        return GF28Utils.toBytes(result);
    }

    /**
     * 密钥扩展
     * @param key
     * @return
     */
    public static byte[] keyExpand(byte[] key) {
        byte[] result = new byte[176];
        // 第0轮
        for (int i = 0; i < 16; i++) result[i] = key[i];
        // 第1~10轮
        for (int j = 1; j <= 10; j++) {
            int count = j * 16 - 1;
            byte[] g = new byte[] {result[count - 12], result[count - 8], result[count - 4], result[count]};
            g = BitsUtils.xor(sBox(ByteUtils.movLeft(g, 1)), new byte[] {rcon(j).toByte(), 0, 0, 0});
            for (int k = 0; k < 4; k++)
                result[j * 16 + 4 * k] = BitsUtils.xor(result[j * 16 - 4 * (4 - k)], g[k]);
            for (int i = 1; i < 4; i++) {
                for (int k = 0; k < 4; k++)
                    result[j * 16 + i + 4 * k] = BitsUtils.xor(result[j * 16 + i - 4 * (4 - k)], result[j * 16 + i - 1 + 4 * k]);
            }
        }
        return result;
    }

    /**
     * 获取已扩展的密钥的一部分
     * @param expandedKey
     * @param round
     * @return
     */
    public static byte[] getExpandedKey(byte[] expandedKey, int round) {
        byte[] result = new byte[16];
        for (int i = 0; i < 16; i++) result[i] = expandedKey[round * 16 + i];
        return result;
    }

    /**
     * 轮常量
     * @param round
     * @return
     */
    public static GF28 rcon(int round) {
        if (round == 1) return GF28.getOneInstance();
        return GF28Utils.multi(new GF28(new byte[]{0, 0, 0, 0, 0, 0, 0, 1, 0}), rcon(round - 1));
    }

    /**
     * 获取乘法逆元
     * @return
     */
    public static byte getInverse(byte b) {
        GF28 input = new GF28(b);
        GF28 m = GF28.M;
        List<GF28> qList = new ArrayList<>();
        List<GF28> rList = new ArrayList<>();
        List<GF28> vList = new ArrayList<>();
        List<GF28> wList = new ArrayList<>();
        rList.add(m);
        rList.add(input);
        vList.add(GF28.getOneInstance());
        vList.add(GF28.getZeroInstance());
        wList.add(GF28.getZeroInstance());
        wList.add(GF28.getOneInstance());
        while (!ByteUtils.isEqual(rList.get(rList.size() - 1).getValues(), GF28.ZERO.getValues())) {
            rList.add(gf28Mod(rList.get(rList.size() - 2), rList.get(rList.size() - 1)));
            qList.add(gf28Div(rList.get(rList.size() - 3), rList.get(rList.size() - 2)));
            vList.add(GF28Utils.xor(vList.get(vList.size() - 2), GF28Utils.multi(qList.get(qList.size() - 1), vList.get(vList.size() - 1))));
            wList.add(GF28Utils.xor(wList.get(wList.size() - 2), GF28Utils.multi(qList.get(qList.size() - 1), wList.get(wList.size() - 1))));
        }
        return wList.get(wList.size() - 2).toByte();
    }

    /**
     * 比较两个位数组的大小
     * b1=b2返回0；b1>b2返回1；b1<b2返回-1
     * @param bits1
     * @param bits2
     * @return
     */
    public static int compareBits(Bits bits1, Bits bits2) {
        for (int i = 0; i < bits1.getLength(); i++) {
            if (bits1.get(i) > bits2.get(i)) return 1;
            else if (bits1.get(i) < bits2.get(i)) return -1;
        }
        return 0;
    }

    /**
     * 在GF(2^8)中取模，即num1 mod num2
     * @param num1
     * @param num2
     * @return
     */
    public static GF28 gf28Mod(GF28 num1, GF28 num2) {
        if (compareBits(num1, num2) == -1) return GF28.getZeroInstance();
        byte[] temp = num1.getValues();
        while (compareBits(new Bits(temp), num2) >= 0) {
            int dis = getHeighestDis(temp, num2.getValues());
            temp = BitsUtils.xor(temp, num2.movLeft(dis).getValues());
        }
        return new GF28(temp);
    }

    /**
     * 在GF(2^8)中相除，即num1|num2的商
     * @param num1
     * @param num2
     * @return
     */
    public static GF28 gf28Div(GF28 num1, GF28 num2) {
        if (compareBits(num1, num2) == -1) return GF28.getZeroInstance();
        byte[] temp = num1.getValues();
        byte divRes = 0;
        while (compareBits(new Bits(temp), num2) >= 0) {
            int dis = getHeighestDis(temp, num2.getValues());
            divRes ^= 1 << dis;
            temp = BitsUtils.xor(temp, num2.movLeft(dis).getValues());
        }
        return new GF28(divRes);
    }


    /**
     * 获取两个位数组的最高位位置之差
     * @param b1
     * @param b2
     * @return
     */
    public static int getHeighestDis(byte[] b1, byte[] b2) {
        int top1 = BitsUtils.getTopPos(b1);
        int top2 = BitsUtils.getTopPos(b2);
        if (top1 == top2) return 0;
        int dis = top1 - top2;
        return dis;
    }

    /**
     * 字节矩阵转为数组
     * @param matrix
     * @return
     */
    public static byte[] byteMatrix2Arr(byte[][] matrix) {
        byte[] arr = new byte[matrix.length * matrix[0].length];
        for (int i = 0; i < 4; i++) {
            for (int j = 0; j < 4; j++) {
                arr[4 * i + j] = matrix[i][j];
            }
        }
        return arr;
    }

    /**
     * 字节数组转为矩阵
     * @param arr
     * @return
     */
    public static byte[][] byteArr2Matrix(byte[] arr) {
        byte[][] matrix = new byte[4][4];
        for (int i = 0; i < 4; i++) {
            for (int j = 0; j < 4; j++) {
                matrix[i][j] = arr[i * 4 + j];
            }
        }
        return matrix;
    }

    /**
     * 将十六进制字符串转换为满足AES要求的字节数组
     * @param str
     * @return
     */
    public static byte[] hexStr2AESBytes(String str) {
        return byteMatrix2Arr(posBytes(byteArr2Matrix(HexUtils.toBytes(str))));
    }

    /**
     * 按照AES要求调整字节数组中的位置
     * @return
     */
    public static byte[][] posBytes(byte[][] input) {
        byte[][] result = new byte[input.length][input[0].length];
        for (int i = 0; i < result.length; i++) {
            for (int j = 0; j < result[0].length; j++) {
                result[j][i] = input[i][j];
            }
        }
        return result;
    }

    /**
     * 行移位变换
     * @param inputRows
     * @return
     */
    public static byte[] movRows(byte[] inputRows) {
        byte[] result = new byte[inputRows.length];
        byte[] row1 = ByteUtils.rangeOf(inputRows, 0, 3);
        byte[] row2 = ByteUtils.rangeOf(inputRows, 4, 7);
        byte[] row3 = ByteUtils.rangeOf(inputRows, 8, 11);
        byte[] row4 = ByteUtils.rangeOf(inputRows, 12, 15);
        movRow(row1, 0);
        movRow(row2, 1);
        movRow(row3, 2);
        movRow(row4, 3);
        for (int i = 0; i < 4; i++) result[i] = row1[i];
        for (int i = 0; i < 4; i++) result[4 + i] = row2[i];
        for (int i = 0; i < 4; i++) result[8 + i] = row3[i];
        for (int i = 0; i < 4; i++) result[12 + i] = row4[i];
        return result;
    }

    /**
     * 单行移位
     * @param inputRow
     * @param movCount
     * @return
     */
    public static byte[] movRow(byte[] inputRow, int movCount) {
        movCount = movCount % 4;
        if (movCount == 0) return inputRow;
        else {
            byte[] temp = new byte[movCount];
            int len = inputRow.length;
            for (int i = 0; i < movCount; i++)
                temp[i] = inputRow[i];
            for (int j = movCount; j < len; j++)
                inputRow[j - movCount] = inputRow[j];
            for (int k = 0; k < movCount; k++)
                inputRow[len - movCount + k] = temp[k];
            return inputRow;
        }
    }

    /**
     * S盒处理
     * @param input
     * @return
     */
    public static byte[] sBox(byte[] input) {
        byte[] result = new byte[input.length];
        for (int i = 0; i < input.length; i++) {
            byte s = input[i];
            result[i] = getSBoxVal(s);
        }
        return result;
    }

    /**
     * 获取一个字节对应S盒的值
     * @param inputByte
     * @return
     */
    public static byte getSBoxVal(byte inputByte) {
        byte[][] m = {
                {1, 0, 0, 0, 1, 1, 1, 1},
                {1, 1, 0, 0, 0, 1, 1, 1},
                {1, 1, 1, 0, 0, 0, 1, 1},
                {1, 1, 1, 1, 0, 0, 0, 1},
                {1, 1, 1, 1, 1, 0, 0, 0},
                {0, 1, 1, 1, 1, 1, 0, 0},
                {0, 0, 1, 1, 1, 1, 1, 0},
                {0, 0, 0, 1, 1, 1, 1, 1}
        };
        byte[] c = {1, 1, 0, 0, 0, 1, 1, 0};
        byte[] result = new byte[8];
        byte[] inputBits = BitsUtils.fromByte(getInverse(inputByte));
        for (int i = 0; i < m.length; i++) {
            byte[] temp = new byte[8];
            for (int j = 0; j < m[i].length; j++)
                temp[j] = (byte) (m[i][j] * inputBits[inputBits.length - 1 - j]);
            byte tempBit = temp[0];
            for (int k = 1; k < temp.length; k++)
                tempBit = BitsUtils.xor(tempBit, temp[k]);
            result[i] = BitsUtils.xor(tempBit, c[i]);
        }
        return BitsUtils.toByte(reverseBits(result));
    }

    /**
     * 将位数组翻转（用于S盒算法）
     * @param input
     * @return
     */
    public static byte[] reverseBits(byte[] input) {
        byte[] result = input;
        byte temp;
        for (int i = 0; i < 4; i++) {
            temp = result[i];
            result[i] = result[7 - i];
            result[7 - i] = temp;
        }
        return result;
    }

    /**
     * 轮密钥加
     * @param pBytes
     * @param kBytes
     * @return
     */
    public static byte[] pkPlus(byte[] pBytes, byte[] kBytes) {
        if (pBytes.length != kBytes.length) return null;
        byte[] result = new byte[pBytes.length];
        for (int i = 0; i < result.length; i++)
            result[i] = BitsUtils.xor(pBytes[i], kBytes[i]);
        return result;
    }

    /**
     * 结果以十六进制字符串格式输出
     * @param des
     * @param bytes
     */
    public static void printHexStrResult(String des, byte[] bytes) {
        System.out.println(des);
        String[] strArr = new String[bytes.length];
        for (int i = 0; i < strArr.length; i++) strArr[i] = ByteUtils.toHex(bytes[i]);
        int temp = 1;
        for (String str: strArr) {
            if (temp <= 3)
                System.out.print(str + " ");
            else {
                System.out.print(str + "\n");
                temp = 0;
            }
            temp++;
        }
    }

}
